package com.kakarote.hbase.config;

import org.apache.hadoop.hbase.Cell;
import org.apache.hadoop.hbase.CellUtil;
import org.apache.hadoop.hbase.TableName;
import org.apache.hadoop.hbase.client.*;
import org.apache.hadoop.hbase.util.Bytes;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;

@Component
public class HBaseTemplate {

    private final Connection connection;

    public HBaseTemplate(Connection connection) {
        this.connection = connection;
    }

    public void createTable(String tableName, String... columnFamilies) throws IOException {
        try (Admin admin = connection.getAdmin()) {
            TableName tn = TableName.valueOf(tableName);
            if (!admin.tableExists(tn)) {
                TableDescriptorBuilder tableDescriptor = TableDescriptorBuilder.newBuilder(tn);
                for (String cf : columnFamilies) {
                    tableDescriptor.setColumnFamily(ColumnFamilyDescriptorBuilder.of(cf));
                }
                admin.createTable(tableDescriptor.build());
            }
        }
    }

    public void putData(String tableName, String rowKey, String columnFamily,
                        Map<String, String> data) throws IOException {
        try (Table table = connection.getTable(TableName.valueOf(tableName))) {
            Put put = new Put(Bytes.toBytes(rowKey));
            for (Map.Entry<String, String> entry : data.entrySet()) {
                put.addColumn(Bytes.toBytes(columnFamily),
                        Bytes.toBytes(entry.getKey()),
                        Bytes.toBytes(entry.getValue()));
            }
            table.put(put);
        }
    }

    public Map<String, String> getData(String tableName, String rowKey,
                                       String columnFamily) throws IOException {
        try (Table table = connection.getTable(TableName.valueOf(tableName))) {
            Get get = new Get(Bytes.toBytes(rowKey));
            Result result = table.get(get);

            Map<String, String> data = new HashMap<>();
            for (Cell cell : result.listCells()) {
                String qualifier = Bytes.toString(CellUtil.cloneQualifier(cell));
                String value = Bytes.toString(CellUtil.cloneValue(cell));
                data.put(qualifier, value);
            }
            return data;
        }
    }
}
